/**
 * Copyright (C) 2021 - 2031 O-Cubes Co., Ltd.
 */

/****************************************************************
 *  @file    sfud.c
 *  @brief   SPI Flash Operation Functions of SFUD
 *  @version v1.0
 *  @date    03. Apr. 2023
 *  @author  gonght
 ****************************************************************/
#include "sfud.h"
#include <string.h>
#include "syscounter.h"
#include "pcl.h"
#include "dw_spi.h"
#include "device_spi.h"

#ifndef SFUD_FLASH_DEVICE_TABLE
	#error "Please configure the flash device information table in (in sfud_cfg.h)."
#endif

#define QSPI_XIP_REG_ADDR                  (MEM_MAP_TOPSC_BASE_ADDR + 0x124)
#define QSPI_XIP_EN                        BIT(0)
#define QSPI_XIP_READADDR_SHIFT_EN         BIT(1)
#define QSPI_XIP_M_ADDR                    (MEM_MAP_TOPSC_BASE_ADDR + 0x128)

/* user configured flash device information table */
static sfud_flash flash_table[] = SFUD_FLASH_DEVICE_TABLE;
/* supported manufacturer information table */
static const sfud_mf mf_table[] = SFUD_MF_TABLE;
/* supported flash chip information table */
static const sfud_flash_chip flash_chip_table[] = SFUD_FLASH_CHIP_TABLE;
uint32_t flash_idx = SFUD_GD25Q40E_DEVICE_INDEX;

extern sfud_err sfud_spi_port_init(sfud_flash *flash);

static sfud_err software_init(const sfud_flash *flash);
static sfud_err hardware_init(sfud_flash *flash);
static sfud_err page256_or_1_byte_write(const sfud_flash *flash, uint32_t addr, size_t size, uint16_t write_gran,
                                        const uint8_t *data);
static sfud_err wait_busy(const sfud_flash *flash);
static sfud_err reset(const sfud_flash *flash);
static sfud_err read_jedec_id(sfud_flash *flash);
static sfud_err set_write_enabled(const sfud_flash *flash, bool enabled);
static void make_adress_byte_array(const sfud_flash *flash, uint32_t addr, uint8_t *array);

/**
 * SFUD initialize by flash device
 *
 * @param flash flash device
 *
 * @return result
 */
static sfud_err sfud_device_init(sfud_flash *flash)
{
	sfud_err result = SFUD_SUCCESS;

	/* hardware initialize */
	result = hardware_init(flash);

	if (result == SFUD_SUCCESS) {
		result = software_init(flash);
	}

	if (result == SFUD_SUCCESS) {
		flash->init_ok = true;
		SFUD_INFO("%s flash device is initialize success.", flash->name);
	} else {
		flash->init_ok = false;
		SFUD_INFO("Error: %s flash device is initialize fail.", flash->name);
	}

	return result;
}

uint8_t sfud_get_mf_id(void)
{
	return flash_chip_table[flash_idx].mf_id;
}

/**
 * SFUD library initialize.
 *
 * @return result
 */
sfud_err sfud_init(void)
{
	sfud_err result = SFUD_SUCCESS;
	size_t i;

	SFUD_DEBUG("SFUD init.");

	/* initialize flash device */
	for (i = 0; i < sizeof(flash_table) / sizeof(sfud_flash); i++) {
		flash_table[i].index = i;
		result = sfud_device_init(&flash_table[i]);

		if (result == SFUD_SUCCESS) {
			break;
		}
	}

	return result;
}

/**
 * hardware initialize
 */
static sfud_err hardware_init(sfud_flash *flash)
{
	sfud_err result = SFUD_SUCCESS;
	uint8_t status[2];

	SFUD_ASSERT(flash);

	result = sfud_spi_port_init(flash);

	if (result != SFUD_SUCCESS) {
		return result;
	}

	/* SPI write read function must be initialize */
	SFUD_ASSERT(flash->spi.wr);

	/* if the user don't configure flash chip information then using SFDP parameter or static flash parameter table */
	if (flash->chip.capacity == 0 || flash->chip.write_mode == 0 || flash->chip.erase_gran == 0
		|| flash->chip.erase_gran_cmd == 0) {
		/* read JEDEC ID include manufacturer ID, memory type ID and flash capacity ID */
		result = read_jedec_id(flash);

		if (result != SFUD_SUCCESS) {
			return result;
		}
	}

	if (flash->chip.capacity == 0 || flash->chip.write_mode == 0 || flash->chip.erase_gran == 0
		|| flash->chip.erase_gran_cmd == 0) {
		SFUD_INFO("Warning: This flash device is not found or not support.");
		return SFUD_ERR_NOT_FOUND;
	} else {
		const char *flash_mf_name = NULL;

		if (mf_table[flash_idx].id == flash->chip.mf_id) {
			flash_mf_name = mf_table[flash_idx].name;
		}

		/* print manufacturer and flash chip name */
		if (flash_mf_name && flash->chip.name) {
			SFUD_INFO("Find a %s %s flash chip. Size is %ld bytes.", flash_mf_name, flash->chip.name,
				flash->chip.capacity);
		} else if (flash_mf_name) {
			SFUD_INFO("Find a %s flash chip. Size is %ld bytes.", flash_mf_name, flash->chip.capacity);
		} else {
			SFUD_INFO("Find a flash chip. Size is %ld bytes.", flash->chip.capacity);
		}
	}

	/* reset flash device */
	result = reset(flash);

	if (result != SFUD_SUCCESS) {
		return result;
	}

	flash->addr_in_4_byte = false;

	if (flash_chip_table[flash_idx].mf_id == SFUD_MF_ID_GIANTEC) {
		sfud_read_status(SFUD_CMD_READ_STATUS_REGISTER2, &status[0], 1);
		status[0] |= BIT(1);
		sfud_write_status(false, SFUD_CMD_WRITE_STATUS_REGISTER2, &status[0], 1);
	} else if (flash_chip_table[flash_idx].mf_id == SFUD_MF_ID_GIGADEVICE) {
		sfud_read_status(SFUD_CMD_READ_STATUS_REGISTER1, &status[0], 1);
		sfud_read_status(SFUD_CMD_READ_STATUS_REGISTER2, &status[1], 1);
		status[1] |= BIT(1);
		sfud_write_status(false, SFUD_CMD_WRITE_STATUS_REGISTER, status, 2);
	}

	return result;
}

/**
 * software initialize
 *
 * @param flash flash device
 *
 * @return result
 */
static sfud_err software_init(const sfud_flash *flash)
{
	sfud_err result = SFUD_SUCCESS;

	SFUD_ASSERT(flash);

	return result;
}

/**
 * read flash data
 *
 * @param addr start address
 * @param size read size
 * @param data read data pointer
 *
 * @return result
 */
sfud_err sfud_read(uint32_t addr, size_t size, uint8_t *data)
{
	sfud_err result = SFUD_SUCCESS;
	sfud_flash *flash = &flash_table[flash_idx];
	const sfud_spi *spi = &flash->spi;
	uint8_t cmd_data[5], cmd_size;

	SFUD_ASSERT(flash);
	SFUD_ASSERT(data);
	/* must be call this function after initialize OK */
	SFUD_ASSERT(flash->init_ok);

	/* check the flash address bound */
	if (addr + size > flash->chip.capacity) {
		SFUD_INFO("Error: Flash address is out of bound.");
		return SFUD_ERR_ADDR_OUT_OF_BOUND;
	}

	/* lock SPI */
	if (spi->lock) {
		spi->lock(spi);
	}

	result = wait_busy(flash);

	if (result == SFUD_SUCCESS) {
		make_adress_byte_array(flash, addr, &cmd_data[1]);
		cmd_size = flash->addr_in_4_byte ? 5 : 4;
#if CONFIG_SDK_QSPI_FLASH_QUAD_MODE
		cmd_data[0] = SFUD_CMD_QUAD_IO_READ_DATA;
		pcl_set_spi3_tf(PCL_SPI_TF_QUAD);
		result = spi->qspi_wr(spi, cmd_data, cmd_size, data, size);
		pcl_set_spi3_tf(PCL_SPI_TF_STD);
#else
		cmd_data[0] = SFUD_CMD_READ_DATA;
		result = spi->wr(spi, cmd_data, cmd_size, data, size);
#endif
	}

	/* unlock SPI */
	if (spi->unlock) {
		spi->unlock(spi);
	}

	return result;
}

/**
 * erase all flash data
 *
 * @return result
 */
sfud_err sfud_chip_erase(void)
{
	sfud_err result = SFUD_SUCCESS;
	sfud_flash *flash = &flash_table[flash_idx];
	const sfud_spi *spi = &flash->spi;
	uint8_t cmd_data[4];

	SFUD_ASSERT(flash);
	/* must be call this function after initialize OK */
	SFUD_ASSERT(flash->init_ok);

	/* lock SPI */
	if (spi->lock) {
		spi->lock(spi);
	}

	/* set the flash write enable */
	result = set_write_enabled(flash, true);

	if (result != SFUD_SUCCESS) {
		goto __exit;
	}

	cmd_data[0] = SFUD_CMD_ERASE_CHIP;

	result = spi->wr(spi, cmd_data, 1, NULL, 0);

	if (result != SFUD_SUCCESS) {
		SFUD_INFO("Error: Flash chip erase SPI communicate error.");
		goto __exit;
	}

	result = wait_busy(flash);

__exit:
	/* set the flash write disable */
	set_write_enabled(flash, false);

	/* unlock SPI */
	if (spi->unlock) {
		spi->unlock(spi);
	}

	return result;
}

/**
 * erase flash data
 *
 * @note It will erase align by erase granularity.
 *
 * @param addr start address
 * @param size erase size
 *
 * @return result
 */
sfud_err sfud_erase(uint32_t addr, size_t size)
{
	sfud_err result = SFUD_SUCCESS;
	sfud_flash *flash = &flash_table[flash_idx];
	const sfud_spi *spi = &flash->spi;
	uint8_t cmd_data[5], cmd_size, cur_erase_cmd;
	size_t cur_erase_size;

	SFUD_ASSERT(flash);
	/* must be call this function after initialize OK */
	SFUD_ASSERT(flash->init_ok);

	/* check the flash address bound */
	if (addr + size > flash->chip.capacity) {
		SFUD_INFO("Error: Flash address is out of bound.");
		return SFUD_ERR_ADDR_OUT_OF_BOUND;
	}

	if (addr == 0 && size == flash->chip.capacity) {
		return sfud_chip_erase();
	}

	/* lock SPI */
	if (spi->lock) {
		spi->lock(spi);
	}

	/* loop erase operate. erase unit is erase granularity */
	while (size) {
		cur_erase_cmd = flash->chip.erase_gran_cmd;
		cur_erase_size = flash->chip.erase_gran;

		/* set the flash write enable */
		result = set_write_enabled(flash, true);

		if (result != SFUD_SUCCESS) {
			goto __exit;
		}

		cmd_data[0] = cur_erase_cmd;
		make_adress_byte_array(flash, addr, &cmd_data[1]);
		cmd_size = flash->addr_in_4_byte ? 5 : 4;
		result = spi->wr(spi, cmd_data, cmd_size, NULL, 0);

		if (result != SFUD_SUCCESS) {
			SFUD_INFO("Error: Flash erase SPI communicate error.");
			goto __exit;
		}

		result = wait_busy(flash);

		if (result != SFUD_SUCCESS) {
			goto __exit;
		}

		/* make erase align and calculate next erase address */
		if (addr % cur_erase_size != 0) {
			if (size > cur_erase_size - (addr % cur_erase_size)) {
				size -= cur_erase_size - (addr % cur_erase_size);
				addr += cur_erase_size - (addr % cur_erase_size);
			} else {
				goto __exit;
			}
		} else {
			if (size > cur_erase_size) {
				size -= cur_erase_size;
				addr += cur_erase_size;
			} else {
				goto __exit;
			}
		}
	}

__exit:
	/* set the flash write disable */
	set_write_enabled(flash, false);

	/* unlock SPI */
	if (spi->unlock) {
		spi->unlock(spi);
	}

	return result;
}

/**
 * write flash data (no erase operate) for write 1 to 256 bytes per page mode or byte write mode
 *
 * @param flash flash device
 * @param addr start address
 * @param size write size
 * @param write_gran write granularity bytes, only support 1 or 256
 * @param data write data
 *
 * @return result
 */
static sfud_err page256_or_1_byte_write(const sfud_flash *flash, uint32_t addr, size_t size, uint16_t write_gran,
						const uint8_t *data)
{
	sfud_err result = SFUD_SUCCESS;
	const sfud_spi *spi = &flash->spi;
	static uint8_t cmd_data[5 + SFUD_WRITE_MAX_PAGE_SIZE];
	uint8_t cmd_size;
	size_t data_size;

	SFUD_ASSERT(flash);
	/* only support 1 or 256 */
	SFUD_ASSERT(write_gran == 1 || write_gran == 256);
	/* must be call this function after initialize OK */
	SFUD_ASSERT(flash->init_ok);

	/* check the flash address bound */
	if (addr + size > flash->chip.capacity) {
		SFUD_INFO("Error: Flash address is out of bound.");
		return SFUD_ERR_ADDR_OUT_OF_BOUND;
	}

	/* lock SPI */
	if (spi->lock) {
		spi->lock(spi);
	}

	/* loop write operate. write unit is write granularity */
	while (size) {
		/* set the flash write enable */
		result = set_write_enabled(flash, true);

		if (result != SFUD_SUCCESS) {
			goto __exit;
		}

		make_adress_byte_array(flash, addr, &cmd_data[1]);
		cmd_size = flash->addr_in_4_byte ? 5 : 4;

		/* make write align and calculate next write address */
		if (addr % write_gran != 0) {
			if (size > write_gran - (addr % write_gran)) {
				data_size = write_gran - (addr % write_gran);
			} else {
				data_size = size;
			}
		} else {
			if (size > write_gran) {
				data_size = write_gran;
			} else {
				data_size = size;
			}
		}

		size -= data_size;
		addr += data_size;

		memcpy(&cmd_data[cmd_size], data, data_size);

#if CONFIG_SDK_QSPI_FLASH_QUAD_MODE
		cmd_data[0] = SFUD_CMD_QUAD_PAGE_PROGRAM;
		pcl_set_spi3_tf(PCL_SPI_TF_QUAD);
		result = spi->qspi_wr(spi, cmd_data, cmd_size + data_size, NULL, 0);
		pcl_set_spi3_tf(PCL_SPI_TF_STD);
#else
		cmd_data[0] = SFUD_CMD_PAGE_PROGRAM;
		result = spi->wr(spi, cmd_data, cmd_size + data_size, NULL, 0);
#endif

		if (result != SFUD_SUCCESS) {
			SFUD_INFO("Error: Flash write SPI communicate error.");
			goto __exit;
		}

		result = wait_busy(flash);

		if (result != SFUD_SUCCESS) {
			goto __exit;
		}

		data += data_size;
	}

__exit:
	/* set the flash write disable */
	set_write_enabled(flash, false);

	/* unlock SPI */
	if (spi->unlock) {
		spi->unlock(spi);
	}

	return result;
}

/**
 * write flash data (no erase operate)
 *
 * @param addr start address
 * @param size write size
 * @param data write data
 *
 * @return result
 */
sfud_err sfud_write(uint32_t addr, size_t size, const uint8_t *data)
{
	sfud_err result = SFUD_SUCCESS;
	sfud_flash *flash = &flash_table[flash_idx];

	if (flash->chip.write_mode & SFUD_WM_PAGE_256B) {
		result = page256_or_1_byte_write(flash, addr, size, 256, data);
	}

	return result;
}

/**
 * erase and write flash data
 *
 * @param addr start address
 * @param size write size
 * @param data write data
 *
 * @return result
 */
sfud_err sfud_erase_write(uint32_t addr, size_t size, const uint8_t *data)
{
	sfud_err result = SFUD_SUCCESS;

	result = sfud_erase(addr, size);

	if (result == SFUD_SUCCESS) {
		result = sfud_write(addr, size, data);
	}

	return result;
}

static sfud_err reset(const sfud_flash *flash)
{
	sfud_err result = SFUD_SUCCESS;
	const sfud_spi *spi = &flash->spi;
	uint8_t cmd_data[2];

	SFUD_ASSERT(flash);

	cmd_data[0] = SFUD_CMD_ENABLE_RESET;
	result = spi->wr(spi, cmd_data, 1, NULL, 0);

	if (result == SFUD_SUCCESS) {
		result = wait_busy(flash);
	} else {
		SFUD_INFO("Error: Flash device reset failed.");
		return result;
	}

	cmd_data[1] = SFUD_CMD_RESET;
	result = spi->wr(spi, &cmd_data[1], 1, NULL, 0);

	if (result == SFUD_SUCCESS) {
		result = wait_busy(flash);
	}

	if (result == SFUD_SUCCESS) {
		SFUD_DEBUG("Flash device reset success.");
	} else {
		SFUD_INFO("Error: Flash device reset failed.");
	}

	return result;
}

static sfud_err read_jedec_id(sfud_flash *flash)
{
	sfud_err result = SFUD_SUCCESS;
	const sfud_spi *spi = &flash->spi;
	uint8_t i = 0;
	uint8_t idx = flash->index;
	uint8_t cmd_data, recv_data[3];

	SFUD_ASSERT(flash);

	for (i = 0; i < 3; i++) {
		cmd_data = SFUD_CMD_JEDEC_ID;
		result = spi->wr(spi, &cmd_data, sizeof(cmd_data), recv_data, sizeof(recv_data));

		if (result == SFUD_SUCCESS) {
			flash->chip.mf_id = recv_data[0];
			flash->chip.type_id = recv_data[1];
			flash->chip.capacity_id = recv_data[2];
			SFUD_DEBUG("The flash device manufacturer ID is 0x%02x, memory type ID is 0x%02x, capacity ID is 0x%02x.",
			           flash->chip.mf_id, flash->chip.type_id, flash->chip.capacity_id);

			if ((flash_chip_table[idx].mf_id == flash->chip.mf_id)
			        && (flash_chip_table[idx].type_id == flash->chip.type_id)
			        && (flash_chip_table[idx].capacity_id == flash->chip.capacity_id)) {
				flash->chip.name = flash_chip_table[idx].name;
				flash->chip.capacity = flash_chip_table[idx].capacity;
				flash->chip.write_mode = flash_chip_table[idx].write_mode;
				flash->chip.erase_gran = flash_chip_table[idx].erase_gran;
				flash->chip.erase_gran_cmd = flash_chip_table[idx].erase_gran_cmd;

				flash_idx = idx;
				break;
			}
		} else {
			SFUD_INFO("Error: Read flash device JEDEC ID error, retry_time = %d.", i);
		}

		syscnt_udelay(1000);
	}

	return result;
}

/**
 * set the flash write enable or write disable
 *
 * @param flash flash device
 * @param enabled true: enable  false: disable
 *
 * @return result
 */
static sfud_err set_write_enabled(const sfud_flash *flash, bool enabled)
{
	sfud_err result = SFUD_SUCCESS;
	uint8_t cmd, register_status;

	SFUD_ASSERT(flash);

	if (enabled) {
		cmd = SFUD_CMD_WRITE_ENABLE;
	} else {
		cmd = SFUD_CMD_WRITE_DISABLE;
	}

	result = flash->spi.wr(&flash->spi, &cmd, 1, NULL, 0);

	if (result == SFUD_SUCCESS) {
		result = sfud_read_status(SFUD_CMD_READ_STATUS_REGISTER1, &register_status, 1);
	}

	if (result == SFUD_SUCCESS) {
		if (enabled && (register_status & SFUD_STATUS_REGISTER_WEL) == 0) {
			SFUD_INFO("Error: Can't enable write status.");
			return SFUD_ERR_WRITE;
		} else if (!enabled && (register_status & SFUD_STATUS_REGISTER_WEL) != 0) {
			SFUD_INFO("Error: Can't disable write status.");
			return SFUD_ERR_WRITE;
		}
	}

	return result;
}

/**
 * read flash register status
 *
 * @param flash flash device
 * @param status register status
 *
 * @return result
 */
sfud_err sfud_read_status(uint8_t cmd_byte, uint8_t *status, uint8_t r_len)
{
	uint8_t cmd;
	sfud_flash *flash = &flash_table[flash_idx];

	SFUD_ASSERT(status);

	cmd = cmd_byte;

	return flash->spi.wr(&flash->spi, &cmd, 1, status, r_len);
}

static sfud_err wait_busy(const sfud_flash *flash)
{
	sfud_err result = SFUD_SUCCESS;
	uint8_t status;
	size_t retry_times = flash->retry.times;

	SFUD_ASSERT(flash);

	while (true) {
		result = sfud_read_status(SFUD_CMD_READ_STATUS_REGISTER1, &status, 1);

		if (result == SFUD_SUCCESS && ((status & SFUD_STATUS_REGISTER_BUSY)) == 0) {
			break;
		}

		/* retry counts */
		SFUD_RETRY_PROCESS(flash->retry.delay, retry_times, result);
	}

	if (result != SFUD_SUCCESS || ((status & SFUD_STATUS_REGISTER_BUSY)) != 0) {
		SFUD_INFO("Error: Flash wait busy has an error.");
	}

	return result;
}

static void make_adress_byte_array(const sfud_flash *flash, uint32_t addr, uint8_t *array)
{
	uint8_t len, i;

	SFUD_ASSERT(flash);
	SFUD_ASSERT(array);

	len = flash->addr_in_4_byte ? 4 : 3;

	for (i = 0; i < len; i++) {
		array[i] = (addr >> ((len - (i + 1)) * 8)) & 0xFF;
	}
}

/**
 * write status register
 *
 * @param flash flash device
 * @param is_volatile true: volatile mode, false: non-volatile mode
 * @param status register status
 *
 * @return result
 */
sfud_err sfud_write_status(bool is_volatile, uint8_t cmd_byte, uint8_t *status, uint8_t w_len)
{
	sfud_err result = SFUD_SUCCESS;
	sfud_flash *flash = &flash_table[flash_idx];
	const sfud_spi *spi = &flash->spi;
	uint8_t cmd_data[w_len + 1];
	uint8_t i = 0;

	if (is_volatile) {
		if (flash->chip.mf_id == SFUD_MF_ID_GIGADEVICE) {
			cmd_data[0] = SFUD_VOLATILE_SR_WRITE_ENABLE;
			result = spi->wr(spi, cmd_data, 1, NULL, 0);
		}
	} else {
		result = set_write_enabled(flash, true);
	}

	if (result == SFUD_SUCCESS) {
		cmd_data[0] = cmd_byte;

		for (i = 0; i < w_len; i++) {
			cmd_data[i + 1] = status[i];
		}

		result = spi->wr(spi, cmd_data, w_len + 1, NULL, 0);
	}

	if (result != SFUD_SUCCESS) {
		SFUD_INFO("Error: Write_status register failed.");
	}

	return result;
}

void sfud_xip_en(void)
{
	uint8_t cont_byte = 0;
	uint8_t read_cmd = 0;
	DEV_SPI_PTR spi_ptr = spi_get_dev(DW_SPI_3_ID);

	read_cmd = SFUD_CMD_QUAD_IO_READ_DATA;

	if (SFUD_MF_ID_GIANTEC == sfud_get_mf_id())
		cont_byte = 0x22;
	else if (SFUD_MF_ID_GIGADEVICE == sfud_get_mf_id())
		cont_byte = 0xAA;

	pcl_set_spi3_tf(PCL_SPI_TF_QUAD);
	dw_spi_en_xip(spi_ptr, read_cmd, cont_byte);
	writel(QSPI_XIP_EN | QSPI_XIP_READADDR_SHIFT_EN, QSPI_XIP_REG_ADDR);
	writel(cont_byte, QSPI_XIP_M_ADDR);
}

void sfud_xip_dis(void)
{
	DEV_SPI_PTR spi_ptr = spi_get_dev(DW_SPI_3_ID);

	writel(0, QSPI_XIP_M_ADDR);
	writel(0, QSPI_XIP_REG_ADDR);
	dw_spi_dis_xip(spi_ptr);
	pcl_set_spi3_tf(PCL_SPI_TF_STD);
	reset(&flash_table[flash_idx]);
}


